
from pocketflow import AsyncNode
from scorpio.core.common import EmbeddingConfig, ModelParameter
from scorpio.core.providers import LLMClient

from scorpio.core.common import get_logger
logger = get_logger(__name__)

# RAG embedding action
class EmbeddingAction(AsyncNode):
    async def prep_async(self, shared):
        """Get the search query from the shared store.
         shared = {
            "context": {
                "embedding_config": self.embedding_config,
                "ingestion": self.ingestion_service
            },
            "document": {
                "file_path": file_path,
                "title": title or os.path.basename(file_path),
            },
            "embedding_config": self.embedding_config
        }
        """
        
        context = shared.get("context")
        document = shared.get("document")
        embedding_config: EmbeddingConfig = context.get("embedding_config")
        model_param: ModelParameter = ModelParameter(
            provider=embedding_config.provider,
            model=embedding_config.model,
            title=embedding_config.model,
            base_url=embedding_config.base_url,
            dimension=embedding_config.dimension,
            api_key="",
        )
        cli: LLMClient= LLMClient(provider=embedding_config.provider, 
                            model_param=model_param) 
        return cli, document 
        
    async def exec_async(self, inputs):
        """Search the web for the given query."""
        # Call the search utility function
        cli, document = inputs
        chunks = document.get("chunk")
        chunk_embeddings = []
        logger.info(f"Start embedding streamable chunks")
        async for chunk in chunks:
            chunk_content = chunk.get("content")
            # to-do: add token count
            chunk_embedding = await cli.client.get_embedding(chunk_content)
            chunk_embeddings.append({
                "content": chunk_content,
                "chunk_index": chunk.get("chunk_index"),
                "metadata": chunk.get("metadata"),
                "embedding": chunk_embedding
            })
        
        logger.info(f"Embedding successfully for {len(chunk_embeddings)} chunks")
        logger.debug(f"First chunk preview: {chunk_embeddings[0]}")
        document["chunk_embeddings"] = chunk_embeddings
        return "embedded"
    
    async def post_async(self, shared, prep_res, exec_res):
        """Save the search results and go back to the decision node."""
        if exec_res != "embedded":
            logger.error(f"failed to embedding chunks") 
        return exec_res